'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-11-19
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

from zhipuai import ZhipuAI
import json


with open("../LLM-config.json", "r") as fp:
    LLM_json = json.load(fp)
    user_key = LLM_json['LLM-key']
    user_chat_model = LLM_json['LLM-name']
    knowledge_base_id = LLM_json['KB-id']

system_prompt = (
    "[Character settings]You are a physicist and English native speaker."
    " Now you need to read the title and the abstract of an article, and"
    " extract keywords for each literature. 你须检索\"\"\"{{knowledge}}\"\"\"以获得更规范的标签\n\n\n "
    "[order]\n1. Keywords are seperated by ','.\n"
    "2.Please search the knowledge base. If there are suitable keywords in the knowledge base, "
        "return the keywords in the knowledge base first.\n"
    "3.If one keyword <kw1> has a parent keyword <kw2>, please return in the format of \"<kw2>/<kw1>\"\n"
    "4.**GIVE ME KEYWORDS ONLY, DO NOT REPEAT TITLE AND ABSTRACT**\n"
    "5.从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找合适的关键词，找不到就自拟。自拟关键词需要加上后缀(ai)\n\n"
)

# system_prompt = ("给你论文的标题和摘要，请告诉我知识库\"\"\"{{knowledge}}\"\"\"中哪几个标签最适合该论文。")


# knowledge_prompt = "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找合适的关键词，并告诉用户所说的关键词是否在\"\"\"{{knowledge}}\"\"\"里\n"

client = ZhipuAI(api_key=user_key) # 请填写您自己的APIKey

result = client.knowledge.used()
print(result)
response = client.chat.completions.create(
    model=user_chat_model,  # 填写需要调用的模型名称
    messages=[
        # {"role": "system", "content": system_prompt},
        {"role": "user", "content": (
            "[Title] Recent Advances in Semiclassical Methods Inspired by Supersymmetric Quantum Mechanics, "
            "[Abstract] Semiclassical methods are essential in analyzing quantum mechanical systems. Although they generally produce approximate results, relatively rare potentials exist for which these methods are exact. Such intriguing potentials serve as crucial test cases for semiclassical approximations. Recent research has demonstrated a deep connection between supersymmetric quantum mechanics and the exactness of semiclassical methods. Specifically, the mathematical form of conventional shape-invariant potentials guarantees exactness in several related situations. In this manuscript, we review these recent results and discuss their significance."
        )},
        # {"role": "user", "content": "non-Hermitian"}
    ],
    tools=[
            {
                "type": "retrieval",
                "retrieval": {
                    "knowledge_id": knowledge_base_id,
                    "prompt_template": system_prompt
                }
            }
    ]
)
print(response.choices[0])

result = client.knowledge.used()
print(result)